#pragma once

#include <vector>
#include <cstddef>
#include <functional>

#include <assert.h>
#include <infiniband/verbs.h>


#define BLOCK_USAGE_RATIO 0.5
#define EXTEND_POOL_SIZE  10 << 30
#define EXTEND_BLOCK_SIZE 64 << 10


//! 函数包装器, 类似函数指针,可理解为是函数类型(C++11引入)
using AllocationCallback = std::function<void(void* ptr, uint32_t lkey, 
                                        uint32_t rkey, int pool_idx)>;


using SimpleAllocationCallback = std::function<void(void* ptr, 
                                uint32_t lkey, uint32_t rkey)>;



//! 内存池
class MemoryPool {
public:
    /**
     *! @brief 内存池构造函数, 分配内存空间并注册MR到指定的PD
     * 
     * @param pool_size   内存池大小
     * @param block_size  内存池块大小
     * @param pd          RDMA保护域
     */
    MemoryPool(size_t pool_size, size_t block_size, struct ibv_pd* pd);



    /**
     *! @brief 析构函数，反注册MR并释放整个内存池空间
     * 
     */
    ~MemoryPool();




    /**
     *! @brief 分配内存块(size必须与内存块大小对齐)
     * 
     * @param size 需要分配内存总大小
     * @param n    内存块数量
     * 
     * @return     返回成功分配的块数量
     */
    int allocate(size_t size, size_t n, SimpleAllocationCallback callback);



    /**
     *! @brief 将特定内存块回收
     * 
     * @param ptr 
     * @param size 
     */
    void deallocate(void* ptr, size_t size);



    //! 获取访问本地内存的lkey
    uint32_t get_lkey() const { return mr_->lkey; }



    //! 获取远程访问该内存的rkey
    uint32_t get_rkey() const { return mr_->rkey; }



    //! 返回该内存池z的总块数
    uint32_t get_total_blocks() const { return total_blocks_; }



    //! 返回该内存池已分配的块数
    uint32_t get_allocated_blocks() const { return allocated_blocks_; }


private:
    void*  pool_;                  // 内存池数据地址
    size_t pool_size_;             // 内存池大小
    size_t block_size_;            // 内存池块大小
    size_t total_blocks_;          // 内存池总块数
    size_t last_search_position_;  // 最后搜索位置
    size_t allocated_blocks_;      // 已分配块数

    // TODO: use judy library to speed up the bitmap?
    std::vector<uint64_t> bitmap_;

    struct ibv_mr* mr_;
    struct ibv_pd* pd_;
};



//! 内存池管理器
class MemoryPoolManager {
public:
    // 根据指定内存池参数创建内存池并放入管理器
    MemoryPoolManager(size_t pool_size, 
    size_t block_size, struct ibv_pd* pd) {
        add_mempool(pool_size, block_size, pd);
    }


    MemoryPoolManager(const MemoryPoolManager& mm) = delete;


    ~MemoryPoolManager() {
        for (auto& pool : mempools_) {
            delete pool;  // 内存池不管理内存地址可直接析构
        }
    }


    /**
     * @brief 使用默认内存池参数创建内存池并加入管理器
     * 
     * @param pd  RDMA保护域
     */
    void add_mempool(struct ibv_pd* pd);



    /**
     * @brief 根据指定参数创建内存池并加入管理器
     * 
     * @param pool_size   内存池大小
     * @param block_size  内存池块大小
     * @param pd          RDMA保护域
     */
    void add_mempool(size_t pool_size, 
        size_t block_size, struct ibv_pd* pd);




    bool allocate(size_t size, size_t n, AllocationCallback callback);




    //! 将特定内存块回收
    void deallocate(void* ptr, size_t size, int pool_idx);




    //! 内存块使用比率
    float usage() {
        size_t total_blocks     = 0;
        size_t allocated_blocks = 0;
        for (auto pool : mempools_) {
            total_blocks += pool->get_total_blocks();
            allocated_blocks += pool->get_allocated_blocks();
        }
        return (float)allocated_blocks / total_blocks;
    }



    //! 获取某个内存池的lkey
    uint32_t get_lkey(int pool_idx) const {
        assert(pool_idx >= 0 && (size_t)pool_idx < mempools_.size());
        return mempools_[pool_idx]->get_lkey();
    }



    //! 获取某个内存池的rkey
    uint32_t get_rkey(int pool_idx) const {
        assert(pool_idx >= 0 && (size_t)pool_idx < mempools_.size());
        return mempools_[pool_idx]->get_rkey();
    }

public:
    bool need_extend = false;


private:
    std::vector<MemoryPool*> mempools_;
};
